import torch


def euclid_dist(features_1, features_2):
    # features_1: 10*60*dim
    # features_2: 10*300*dim

    batch_size = features_1.shape[0]  # 10
    dim = features_1.shape[2]  # 12544
    n_way = features_1.shape[1]  # 5
    total_query = features_2.shape[1]  # 300

    # features_1: 10*60*dim -> 10*1*60*dim -> 10*300*60*dim
    features_1 = features_1.view(batch_size, 1, n_way, dim).expand(batch_size, total_query, n_way, dim)
    # features_2: 10*300*dim -> 10*300*1*dim -> 10*300*60*dim
    features_2 = features_2.view(batch_size, total_query, 1, dim).expand(batch_size, total_query, n_way, dim)
    # sim_matrix: 10*300*60
    sim_matrix = torch.sum((features_1 - features_2) ** 2, dim=3)

    sim_matrix = -sim_matrix

    return sim_matrix


# @SIMILARITY.register('regions_euclid')
# def regions_euclid_dist(s_region_f, q_region_f):
#     # features_1: 10*5*2*dim
#     # features_2: 10*75*2*dim
#     pass
